#!/usr/bin/python

import rospy
from yaml import dump
from geometry_msgs.msg import PoseArray, Pose   
#  Pose
#  在自由空间中姿态的代表，包括位置和朝向
#  posesPoint position Quaternion orientation

#  PoseArray: Header header Pose[] 是Pose的列表

from detection_clustering import DetectionClustering 

class DetectionCollector(object):
    def __init__(self):  
        # 构造函数 
        #初始化 
        #  1： 订阅 检测的 pose 方位 话题；
        #  2： 检测类别；

        self.detected = {}
        self.detection_names = rospy.get_param('/darknet_ros/yolo_model/detection_classes/names')
        rospy.Subscriber('/cluster_decomposer/centroid_pose_array', PoseArray, self.collect)  # 订阅目标的 pose列表 话题
        print 'Searching for objects...'
        
    def __del__(self):
        # 析构函数
        print "保存数据到数据库 detections_raw.db..."  # 保存数据库文件
        with open('detections_raw.db', 'w') as outfile:
            dump(self.detected, outfile) # 包含位置xyz、识别类别

        print "保存点云数据到数据库 detections_dbscan.db..."
        dc = DetectionClustering(self.detected)  
        # 保存点云数据
        with open('detections_dbscan.db', 'w') as outfile:
            dump(dc.clusters, outfile)

    def update_key(self, key, val):
        # 添加数据到 检测话题列表 等待被添加到数据库
        if self.detected.has_key(key):
            self.detected[key].append(val)
        else:
            self.detected[key] = [val]

    # 回调函数 pose
    def collect(self, msg):
        # 遍历检测的pose话题
        for i, pose in enumerate(msg.poses):
            if pose != Pose():
                pos = pose.position
                val = [pos.x, pos.y, pos.z] #位置
                key = self.detection_names[i] #
                print 'Found a {} at {}'.format(key, val)
                self.update_key(key, val)

if __name__ == '__main__':
    rospy.init_node('detection_collector')
    DetectionCollector()
    rospy.spin()
